import json
from datetime import datetime
from enum import Enum
from typing import Optional, List

from pydantic import field_validator
from sqlalchemy import Column, Text, CHAR, JSON, DateTime, text, UniqueConstraint, BigInteger, \
    delete, update, select, and_, func
from sqlmodel import SQLModel, Field

from app.api.database.base import session_getter, SQLModelSerializable


class LLMServerBase(SQLModelSerializable):
    name: str = Field(default='', index=True, unique=True, description='服务名称',
                      sa_column_kwargs={'comment': '服务名称'})
    description: Optional[str] = Field(default='', sa_column=Column(Text, comment='服务描述'), description='服务描述')
    server_type: str = Field(sa_column=Column(CHAR(20), comment='服务提供方类型'), description='服务提供方类型', )
    # server_api_type: str = Field(sa_column=Column(CHAR(20), comment='服务提供方接口类型'), description='服务提供方类型', )
    config: Optional[dict] = Field(sa_column=Column(JSON, comment='服务提供方公共配置'),
                                   description='服务提供方公共配置')
    user_id: int = Field(default=0, description='创建人ID', sa_type=BigInteger,
                         sa_column_kwargs={'comment': '创建该模型记录的用户ID', })
    create_time: Optional[datetime] = Field(default=None, sa_column=Column(
        DateTime, nullable=False, index=True, server_default=text('CURRENT_TIMESTAMP'), comment='记录创建时间'))
    update_time: Optional[datetime] = Field(default=None, sa_column=Column(
        DateTime, nullable=False, server_default=text('CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP'),
        comment='记录最后更新时间'))


class LLMModelBase(SQLModelSerializable):
    server_id: Optional[int] = Field(default=None, nullable=False, sa_type=BigInteger, index=True, description='服务ID',
                                     sa_column_kwargs={'comment': '服务ID'})
    name: str = Field(default='', description='模型展示名', sa_column_kwargs={'comment': '模型展示名'})
    description: Optional[str] = Field(default='', sa_column=Column(Text, comment='模型描述'), description='模型描述')
    model_name: str = Field(default='', description='模型名称，实例化组件时用的参数',
                            sa_column_kwargs={'comment': '模型名称，实例化组件时用的参数'})
    model_type: str = Field(sa_column=Column(CHAR(20), comment='模型类型'),
                            description='模型类型:  llm, embedding,rerank')
    model_category: Optional[List[str]] = Field(sa_column=Column(JSON, comment='模型分类'),
                                          description='模型分类:  视觉,自然语言处理,语音,其他')
    model_keys: Optional[List[str]] = Field(sa_column=Column(JSON, comment='模型密钥列表'),
                                      description='模型密钥列表'
                                      )
    config: Optional[dict] = Field(default_factory=dict, sa_column=Column(JSON, comment='模型配置'),
                                   description='服务提供方公共配置')
    status: int = Field(default=2, description='模型状态。0：正常，1：异常, 2: 未知',
                        sa_column_kwargs={'comment': '模型状态。0：正常，1：异常, 2: 未知'})
    remark: Optional[str] = Field(default='', sa_column=Column(Text, comment='异常原因'), description='异常原因')
    online: bool = Field(default=True, description='是否在线',
                         sa_column_kwargs={'comment': '是否启用该模型：1=在线，0=下线'})
    user_id: int = Field(default=0, description='创建人ID', sa_type=BigInteger,
                         sa_column_kwargs={'comment': '创建该模型记录的用户ID'})
    create_time: Optional[datetime] = Field(default=None, sa_column=Column(
        DateTime, nullable=False, index=True, server_default=text('CURRENT_TIMESTAMP'), comment='记录创建时间'))
    update_time: Optional[datetime] = Field(default=None, sa_column=Column(
        DateTime, nullable=False, server_default=text('CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP'),
        comment='记录最后更新时间'))



class LLMServer(LLMServerBase, table=True):
    __tablename__ = 'llm_server'

    id: Optional[int] = Field(default=None, sa_type=BigInteger, nullable=False, primary_key=True,
                              description='服务唯一ID',
                              sa_column_kwargs={'comment': '服务唯一ID'})


class LLMModel(LLMModelBase, table=True):
    __tablename__ = 'llm_model'
    __table_args__ = (UniqueConstraint('server_id', 'model_name', name='server_model_uniq'),)

    id: Optional[int] = Field(default=None, sa_type=BigInteger, nullable=False, primary_key=True,
                              description='模型唯一ID',
                              sa_column_kwargs={'comment': '模型唯一ID'})


class LLMServerType(Enum):
    OPENAI = 'openai'
    ALI_BAILIAN = 'ali_bailian'
    DEEPSEEK = 'deepseek'
    SILICONFLOW = 'siliconflow'
    BAIDU_QIANFAN='baidu_qianfan'
    MOONSHOT =  'moonshot'
    ZHIPU = 'zhipu'
    OLLAMA = 'ollama'
    VOLCENGINE = 'volcengine'


class LLMModelType(Enum):
    LLM = 'llm'
    EMBEDDING = 'embedding'
    RERANK = 'rerank'


class LLMModelCategory(Enum):
    TEXT = 'text'
    EMBEDDING = 'embedding'
    RERANK = 'rerank'
    AUDIO = 'audio'
    VISUAL = 'visual'


class LLMModelInfo(LLMModelBase):
    id: Optional[int] = None


class LLMServerInfo(LLMServerBase):
    id: Optional[int] = None
    models: List[LLMModelInfo] = Field(default_factory=list, description='模型列表')


class LLMDao:

    @classmethod
    def get_server_by_name(cls, name: str):
        """根据服务商名称获取服务商"""
        with session_getter() as session:
            statement = select(LLMServer).where(and_(LLMServer.name == name))
            return session.exec(statement).first()

    @classmethod
    def add_server(cls, server: LLMServer):
        """添加服务商"""
        with session_getter() as session:
            session.add(server)
            session.commit()
            session.refresh(server)
            return server

    @classmethod
    def update_server(cls, server: LLMServer):
        """更新服务商"""
        with session_getter() as session:
            # 存在主键会执行UPDATE
            session.add(server)
            session.commit()
            session.refresh(server)
            return server

    @classmethod
    def delete_server(cls, servers: list[LLMServer]):
        """批量删除服务商"""
        with session_getter() as session:
            session.exec(
                delete(LLMServer).where(LLMServer.id.in_([server.id for server in servers])))
            session.commit()
            return servers

    @classmethod
    def get_servers(cls, user_id: BigInteger, page: int, limit: int):
        with session_getter() as session:
            statement = select(LLMServer).where(and_(LLMServer.user_id == user_id))
            count_statement = session.query(func.count(
                LLMServer.id)).where(and_(LLMServer.user_id == user_id))
            statement = statement.offset(
                (page - 1) * limit
            ).limit(limit).order_by(
                LLMServer.update_time.desc()
            )
            return session.exec(statement).all(), session.exec(count_statement).scalar()

    @classmethod
    def get_server_by_id(cls, server_id: int) ->  Optional[LLMServer]:
        """获取服务商"""
        with session_getter() as session:
            statement = select(LLMServer).where(and_(LLMServer.id == server_id))
            return session.exec(statement).scalar_one_or_none()

    @classmethod
    def add_model(cls, model: LLMModel):
        """添加服务商下的模型"""
        with session_getter() as session:
            session.add(model)
            session.commit()
            session.refresh(model)
            return model

    @classmethod
    def update_models(cls, server: LLMServer, models: list[LLMModel]):
        """批量添加服务商下的模型"""
        with session_getter() as session:
            add_models = []
            update_models = []
            for model in models:
                if model.id:
                    update_models.append(model)
                else:
                    add_models.append(model)
            # 删除模型
            session.exec(
                delete(LLMModel).where(LLMModel.server_id == server.id,
                                       LLMModel.id.not_in([model.id for model in update_models])))
            # 添加新增的模型
            session.add_all(add_models)
            # 更新已有模型的数据
            for one in update_models:
                session.exec(
                    update(LLMModel).where(LLMModel.id == one.id).values(
                        name=one.name,
                        description=one.description,
                        model_name=one.model_name,
                        model_type=one.model_type,
                        model_category=one.model_category,
                        model_keys=one.model_keys,
                        config=one.config))

            session.commit()
            return models

    @classmethod
    def get_server_models(cls, server_id: int):
        """获取服务商下的模型"""
        with session_getter() as session:
            statement = select(LLMModel).where(and_(LLMModel.server_id == server_id))
            return session.exec(statement).all()

    @classmethod
    def get_model_by_id(cls, model_id: int) -> Optional[LLMModel]:
        """获取模型"""
        with session_getter() as session:
            statement = select(LLMModel).where(and_(LLMModel.id == model_id))
            return session.exec(statement).scalar_one_or_none()
